import os
from business.utils.document_loader import DocumentLoader
from business.utils.vector_store import VectorStore
from config.index import Config
import sys

def list_markdown_files(directory):
    """递归列出目录下所有的 markdown 文件"""
    markdown_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.md'):
                full_path = os.path.join(root, file)
                markdown_files.append(full_path)
    return markdown_files

def main():
    try:
        # 获取配置
        cfg = Config()
        
        # 获取当前工作目录
        current_dir = os.getcwd()
        print(f"\n当前工作目录: {current_dir}")
        
        # 设置文件路径
        docs_dir = os.path.join(current_dir, "data/docs/p9app")
        vector_store_path = os.path.join(current_dir, "data/vector")
        
        print(f"文档目录路径: {docs_dir}")
        print(f"向量库路径: {vector_store_path}")
        print(f"系统默认编码：{sys.getfilesystemencoding()}")
        
        # 先列出所有 markdown 文件
        print("\n扫描目录中的所有 Markdown 文件:")
        markdown_files = list_markdown_files(docs_dir)
        print(f"找到 {len(markdown_files)} 个 Markdown 文件:")
        for file in markdown_files:
            print(f"- {os.path.relpath(file, docs_dir)}")
        
        # 确保目录存在并检查权限
        for path in [docs_dir, vector_store_path]:
            try:
                os.makedirs(path, exist_ok=True)
                print(f"\n目录创建/确认成功: {path}")
                # 检查目录权限
                test_file = os.path.join(path, "test.txt")
                with open(test_file, 'w') as f:
                    f.write("test")
                os.remove(test_file)
                print(f"目录权限检查通过: {path}")
            except Exception as e:
                print(f"目录操作失败 {path}: {str(e)}")
                return
        
        # 1. 加载文档目录下的所有文件
        print("\n正在加载文档...")
        loader = DocumentLoader()
        documents = loader.load_directory(
            docs_dir, 
            supported_extensions=['.md']  # 只加载 markdown 文件
        )
        
        if not documents:
            print(f"警告：在目录 {docs_dir} 中没有找到任何文档")
            print("请检查 DocumentLoader 的实现是否支持递归扫描子目录")
            return
            
        total_docs = len(documents)
        total_chunks = sum(len(texts) for texts in documents.values())
        print(f"成功加载 {total_docs} 个文档，共 {total_chunks} 个文本段落")
        
        # 打印每个加载的文档
        print("\n已加载的文档:")
        for filename in documents.keys():
            print(f"- {filename}")

        # 2. 创建/更新向量存储
        print("\n正在处理向量存储...")
        vector_store = VectorStore()
        
        # 合并所有文档的文本
        all_texts = []
        for filename, texts in documents.items():
            print(f"处理文档: {filename}")
            texts_with_source = [f"来源文件: {filename}\n\n{text}" for text in texts]
            all_texts.extend(texts_with_source)
        
        print(f"\n总共处理了 {len(all_texts)} 段文本")
        
        # 增量更新向量存储
        print("\n开始创建向量存储...")
        vs = vector_store.create_vector_store(all_texts, vector_store_path)
        print("向量存储创建完成")
        
        # 3. 测试相似性搜索
        print("\n测试相似性搜索:")
        query = "P9APP是什么？"
        results = vs.similarity_search(query, k=3)
        
        print("\n搜索结果:")
        for i, doc in enumerate(results, 1):
            print(f"\n--- 结果 {i} ---")
            print(doc.page_content)

    except Exception as e:
        print(f"\n发生错误: {e}")
        print(f"错误类型: {type(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 